import sys, os
sys.path.append(f"{os.path.dirname(__file__)}")
sys.path.append(f"{os.path.dirname(__file__)}/../repESP_old")

import cube_helpers, charges, resp_helpers
from rep_esp2 import calc_grid_field, calc_non_grid_field

import argparse
import charges_parser

def main():

    help_description = """
        Reproduce ESP from given charges and save as a cube or .esp file.

        If you do not want to process the charge in any way, raw charges can be
        extracted from Gaussian .log or AIMAll .sumviz files. However, it is
        probably sensible to average or equivalence the charges first, using the
        'average.py' script from this suite. The charges can then be read in from
        the output text file.
        """

    parser = argparse.ArgumentParser(
        parents=[charges_parser.parser],
        description=help_description,
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)

    parser.add_argument("template_file",
                        help="""this file specifies the points at which the ESP is
                        to be reproduced. It can be either a Gaussian cube file
                        (.cub) or a file with ESP fitting points (.esp).

                        In the case of the cube file the content doesn't matter,
                        only the grid and atomic coordinates are used. An .esp file
                        can be generated by running a Gaussian Pop calculation (for
                        a selected ESP-based method) with IOp(6/50)=1.""",
                        metavar="TEMPLATE_FILE")

    parser.add_argument("-o", "--output",
                        help="output file path and name without extension",
                        default='rep_esp')

    args = parser.parse_args()

    filetype = args.template_file[-4:]
    if filetype not in ['.cub', '.esp']:
        print("The template file must have the extension .cub or .esp")
        sys.exit(1)

    if os.path.exists(args.output + filetype):
        raise FileExistsError("Output file exists: " + args.output + filetype)

    if filetype == '.cub':
        template_cube = cube_helpers.Cube(args.template_file)
        molecule = template_cube.molecule
    elif filetype == '.esp':
        esp_file = resp_helpers.G09_esp(args.template_file)
        molecule = esp_file.molecule

    input_type = charges_parser.input_type(args.input_charge_type)
    charges._get_charges(args.input_charge_type, args.input_charge_file,
                         input_type, molecule)

    if filetype == '.cub':
        rep_field = calc_grid_field(molecule, template_cube.field.grid, 'rep_esp',
                                    [args.input_charge_type])[0]

        rep_field.write_cube(args.output + '.cub', molecule,
                             args.input_charge_type)
    elif filetype == '.esp':
        rep_field = calc_non_grid_field(molecule, esp_file.field.points, 'rep_esp',
                                        [args.input_charge_type])[0]
        rep_field.write_to_file(args.output + '.esp', molecule)
